import { WebSocketServer } from 'ws';
import { MCPServer } from './mcp-server';
import { SpreadJSAgent } from './agent';
import { AGENT_PORT } from './config';

interface ClientMessage {
    type: 'user_request' | 'execution_result';
    data: {
        message?: string;
        sessionId: string;
        executionId?: string;
        result?: any;
    };
}

class SpreadJSAgentServer {
    private wss: WebSocketServer;
    private agent: SpreadJSAgent;

    constructor() {
        this.wss = new WebSocketServer({ port: AGENT_PORT });

        // Initialize Agent (will connect to MCP Server via HTTP)
        this.agent = new SpreadJSAgent();

        this.setupWebSocket();
        this.initializeAgent();
    }

        private async initializeAgent(): Promise<void> {
        try {
            // Wait a moment for MCP Server to start up
            await new Promise(resolve => setTimeout(resolve, 1000));
            await this.agent.initialize();
            console.log('[Agent Server] Agent initialized successfully');
        } catch (error) {
            console.error('[Agent Server] Failed to initialize agent:', error);
        }
    }

    private setupWebSocket(): void {
        this.wss.on('connection', (ws) => {
            console.log('Client connected');

            ws.on('message', (message: string) => {
                (async () => {
                    try {
                        const clientMessage: ClientMessage = JSON.parse(message);

                        if (clientMessage.type === 'user_request' && clientMessage.data.message) {
                            await this.handleUserRequest(ws, {
                                message: clientMessage.data.message,
                                sessionId: clientMessage.data.sessionId
                            });
                        } else if (clientMessage.type === 'execution_result') {
                            await this.handleExecutionResult(clientMessage.data);
                        }
                    } catch (error) {
                        this.sendMessage(ws, {
                            type: 'error',
                            data: { message: 'Invalid message format' }
                        });
                    }
                })();
            });

            ws.on('close', () => {
                console.log('Client disconnected');
            });
        });

        console.log(`[Agent Server] WebSocket server running on port ${AGENT_PORT}`);
    }

    private async handleUserRequest(ws: any, data: { message: string; sessionId: string }): Promise<void> {
        try {
            // Execute agent with streaming updates
            const stream = this.agent.executeStream(data.message, data.sessionId);

            for await (const step of stream) {
                // step 已经是完整的消息对象，直接发送
                this.sendMessage(ws, step);
            }

            this.sendMessage(ws, {
                type: 'agent_complete',
                data: { sessionId: data.sessionId }
            });
        } catch (error) {
            this.sendMessage(ws, {
                type: 'error',
                data: { message: error instanceof Error ? error.message : 'Unknown error' }
            });
        }
    }

        private async handleExecutionResult(data: { sessionId: string; executionId?: string; result?: any }): Promise<void> {
        console.log('[Agent Server] Received execution result:', {
            sessionId: data.sessionId,
            executionId: data.executionId,
            result: data.result
        });
    }

    private sendMessage(ws: any, message: any): void {
        ws.send(JSON.stringify(message));
    }

    public async handleSpreadJSQuery(query: string): Promise<any> {
        // This will be called from the frontend when executing SpreadJS queries
        return new Promise((resolve) => {
            // Simulate async execution
            setTimeout(() => {
                resolve({ success: true, result: `Query executed: ${query}` });
            }, 100);
        });
    }

    // Get Agent instance for external access
    public getAgent(): SpreadJSAgent {
        return this.agent;
    }
}

// Start MCP Server on port 3002
console.log('Starting MCP Server...');
const mcpServer = new MCPServer();
mcpServer.start();

// Start Agent Server on port 3001
console.log('Starting Agent Server...');
const agentServer = new SpreadJSAgentServer();

// Export for potential external access
export { mcpServer, agentServer };
